package main

import (
	"context"
	"flag"
	"fmt"

	"net/http"
	"os"

	"strconv"
	"strings"
	"time"

	"gitee.com/kvenli/dm_exporter/collector"
	"gitee.com/kvenli/dm_exporter/common"

	"log/slog"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	timeoutOffset float64
	//metricsPath   string = "test"
	//version       string = "v1.0"
	listenAddress string
	help          bool
	disable       string //命令行传入的需要关闭的指标
	//disables      []string //处理命令行传入的根据,分割为一个切片做处理
	config string //配置文件路径
	c      = common.DMConfigHandler{
		Config: &common.Config{},
	}
)

func init() {
	flag.StringVar(&config, "config-file", "dm.yaml", "Path to the configuration file.")
	flag.StringVar(&listenAddress, "addr", ":9105", "Address to listen on for web interface and telemetry.")
	flag.BoolVar(&help, "h", false, "Show help.")
	flag.StringVar(&disable, "disable", "", "Comma-separated list of collectors to disable.")
}

// 筛选采集指标
func filterScrapers(scrapers []collector.Scraper, collectParams []string) []collector.Scraper {
	var filteredScrapers []collector.Scraper

	// Check if we have some "collect[]" query parameters.
	if len(collectParams) > 0 {
		filters := make(map[string]bool)
		for _, param := range collectParams {
			filters[param] = true
		}

		for _, scraper := range scrapers {
			if filters[scraper.Name()] {
				filteredScrapers = append(filteredScrapers, scraper)
			}
		}
	}
	if len(filteredScrapers) == 0 {
		return scrapers
	}
	return filteredScrapers
}

func getScrapeContext(r *http.Request, logger *slog.Logger) (context.Context, string, string, string, error) {
	params := r.URL.Query()
	target := params.Get("target")
	backupPath := params.Get("backup_path")

	authModule := params.Get("auth_module")
	if authModule == "" {
		authModule = "client"
	}

	cfg := c.GetConfig()
	cfgsection, ok := cfg.Sections[authModule]
	if !ok {
		return nil, "", "", "", fmt.Errorf("could not find section [%s] from config file", authModule)
	}

	dsn, err := cfgsection.FormDSN(target)
	if err != nil {
		return nil, "", "", "", fmt.Errorf("error forming dsn from config section [%s]: %s", authModule, err)
	}

	// If backup_path is not provided in the URL, try to get it from the config.
	if backupPath == "" {
		backupPath = cfgsection.BackupPath
	}

	ctx := r.Context()
	if v := r.Header.Get("X-Prometheus-Scrape-Timeout-Seconds"); v != "" {
		timeoutSeconds, err := strconv.ParseFloat(v, 64)
		if err != nil {
			logger.Error("Failed to parse timeout from Prometheus header", "error", err)
		} else {
			if timeoutOffset >= timeoutSeconds {
				logger.Error("Timeout offset should be lower than prometheus scrape timeout", "offset", timeoutOffset, "prometheus_scrape_timeout", timeoutSeconds)
			} else {
				timeoutSeconds -= timeoutOffset
			}
			var cancel context.CancelFunc
			ctx, cancel = context.WithTimeout(ctx, time.Duration(timeoutSeconds*float64(time.Second)))
			defer cancel()
		}
	}

	return ctx, dsn, target, backupPath, nil
}

// 创建httpHandlerFunc
func newHandler(metrics collector.Metrics, scrapers []collector.Scraper, logger *slog.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		cfg := c.GetConfig()
		cfgsection, ok := cfg.Sections["client"]
		if !ok {
			logger.Error("Failed to parse section [client] from config file")
			http.Error(w, "Failed to parse section [client] from config file", http.StatusInternalServerError)
			return
		}
		dsn, err := cfgsection.FormDSN("")
		if err != nil {
			logger.Error("Failed to form dsn from section [client]", "error", err)
			http.Error(w, "Failed to form dsn from section [client]", http.StatusInternalServerError)
			return
		}

		ctx := r.Context()
		// ... (timeout logic remains the same)

		collect := r.URL.Query()["collect[]"]
		filteredScrapers := filterScrapers(scrapers, collect)

		registry := prometheus.NewRegistry()
		registry.MustRegister(collector.New(ctx, dsn, cfgsection.BackupPath, metrics, filteredScrapers, logger))

		gatherers := prometheus.Gatherers{
			registry,
		}
		h := promhttp.HandlerFor(gatherers, promhttp.HandlerOpts{})
		h.ServeHTTP(w, r)
	}
}

func handleProbe(metrics collector.Metrics, scrapers []collector.Scraper, logger *slog.Logger) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		target := r.URL.Query().Get("target")
		if target == "" {
			http.Error(w, "Target parameter is missing", http.StatusBadRequest)
			return
		}

		ctx, dsn, _, backupPath, err := getScrapeContext(r, logger)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		collectParams := r.URL.Query()["collect[]"]
		filteredScrapers := filterScrapers(scrapers, collectParams)

		registry := prometheus.NewRegistry()
		registry.MustRegister(collector.New(ctx, dsn, backupPath, metrics, filteredScrapers, logger))

		h := promhttp.HandlerFor(registry, promhttp.HandlerOpts{})
		h.ServeHTTP(w, r)
	}
}
func main() {

	// Parse flags.
	flag.Parse()
	if help {
		flag.Usage()
		return
	}

	logger := common.Log
	//加载配置文件
	c.ReloadConfig(config, logger)

	// 从命令行标志和环境变量中获取要禁用的采集器
	disabledScrapers := make(map[string]bool)
	disableCollectors := os.Getenv("DM_EXPORTER_DISABLE_COLLECTORS")
	if disable != "" {
		if disableCollectors != "" {
			disableCollectors = disableCollectors + "," + disable
		} else {
			disableCollectors = disable
		}
	}
	if disableCollectors != "" {
		for _, s := range strings.Split(disableCollectors, ",") {
			disabledScrapers[strings.TrimSpace(s)] = true
		}
	}

	// 获取所有已注册的采集器并启用它们
	enabledScrapers := []collector.Scraper{}
	for name, scraper := range collector.GetScrapers() {
		if !disabledScrapers[name] {
			logger.Info("Scraper enabled", "scraper", name)
			enabledScrapers = append(enabledScrapers, scraper)
		} else {
			logger.Info("Scraper disabled", "scraper", name)
		}
	}

	handlerFunc := newHandler(collector.NewMetrics(), enabledScrapers, logger)
	http.Handle("/metrics", promhttp.InstrumentMetricHandler(prometheus.DefaultRegisterer, handlerFunc))
	http.HandleFunc("/probe", handleProbe(collector.NewMetrics(), enabledScrapers, logger))
	http.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK"))
	})
	//监听端口
	if err := http.ListenAndServe(listenAddress, nil); err != nil {
		logger.Error("Error occur when start server", "error", err)
	}
}
